package com.xb.ai.repository;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.xb.ai.entity.po.Chat;
import com.xb.ai.entity.po.Msg;
import com.xb.ai.service.IChatService;
import jakarta.annotation.Resource;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

@Component
public class InMysqlChatMemory implements ChatMemory {

    @Resource
    private IChatService chatService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        Chat chat = chatService.lambdaQuery()
                .eq(Chat::getChatId, Long.valueOf(conversationId))
                .one();
        if (chat == null) {
            chat = new Chat();
            chat.setChatId(Long.valueOf(conversationId));
            List<Msg> msgList = messages.stream()
                    .map(Msg::new)
                    .toList();
            String messagesJson = JSON.toJSONString(msgList, SerializerFeature.DisableCircularReferenceDetect);
            chat.setMessage(messagesJson);
            chatService.save(chat);
        } else {
            String message = chat.getMessage();
            String messagesJson;
            if (StringUtils.isNotBlank(message) && !message.equals("null")){
                List<Msg> msgList = JSON.parseArray(message, Msg.class);
                msgList.addAll(messages.stream().map(Msg::new).toList());
                messagesJson = JSON.toJSONString(msgList, SerializerFeature.DisableCircularReferenceDetect);
            } else {
                List<Msg> msgList = messages.stream()
                        .map(Msg::new)
                        .toList();
                messagesJson = JSON.toJSONString(msgList, SerializerFeature.DisableCircularReferenceDetect);
            }
            chatService.lambdaUpdate()
                    .eq(Chat::getId, chat.getId())
                    .eq(Chat::getChatId, chat.getChatId())
                    .set(Chat::getMessage, messagesJson)
                    .update();
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        Long chatId = Long.valueOf(conversationId);
        Chat chat = chatService.lambdaQuery()
                .eq(Chat::getChatId, chatId)
                .one();
        if (chat == null) {
            return new ArrayList<>();
        }
        String message = chat.getMessage();
        if (StringUtils.isBlank(message)){
            return new ArrayList<>();
        }
        List<Msg> msgList = JSON.parseArray(message, Msg.class);
        List<Message> messageList = msgList.stream()
                .map(Msg::toMessage)
                .toList();
        int size = messageList.size();
        if (lastN >= size) {
            return messageList;
        } else {
            return messageList.subList(size - lastN - 1, size);
        }
    }

    @Override
    public void clear(String conversationId) {
        chatService.lambdaUpdate()
                .eq(Chat::getChatId, Long.valueOf(conversationId))
                .remove();
    }
}
